/*
 * KiWin32K.cpp
 *
 *  Created on: 2013-10-18
 *      Author: WinDDK
 */

#include "KiWin32K.h"

KiWin32K::KiWin32K()
{
    RtlZeroMemory(mKernelName, 0, KERNELNAMELENGTH + 1);
    mKernelImageBase = 0;
    KiDebug("Success");
}

KiWin32K::~KiWin32K()
{
    KiDebug("Success");
}

BOOLEAN KiWin32K::SetKernelInfo(PCHAR kernelName, ULONG kernelImageBase)
{
    if (!MmIsAddressValid((PVOID) kernelImageBase))
    {
        return FALSE;
    }
    strncpy(mKernelName, "ntoskrnl.exe", KERNELNAMELENGTH);
    mKernelImageBase = kernelImageBase;
    KiDebug("Success");
    return TRUE;
}

BOOLEAN KiWin32K::GetModulePath()
{
    strcpy(mModulePath, "\\??\\");
    if (!GetSystemPath(this->mSystemPath))
    {
        KiDebug("GetSystemPath Failed");
        return FALSE;
    }

    strcpy(this->mModuleName, "win32k.sys");
    KiDebug("%s", this->mModuleName);

    strcat(mModulePath, mSystemPath);
    strcat(mModulePath, "\\");
    strcat(mModulePath, mModuleName);
    KiDebug("%s", mModulePath);
    return TRUE;
}

BOOLEAN KiWin32K::FixModule()
{
    if (!FixRelocTable())
    {
        KiDebug("FixModuleRelocTable Failed");
        return FALSE;
    }

    if (!FixImportTable())
    {
        KiDebug("FixModuleImportTable Failed");
        return FALSE;
    }

    return TRUE;
}

BOOLEAN KiWin32K::FixImportTable()
{
    PIMAGE_IMPORT_DESCRIPTOR imageImportDescriptor;
    ULONG importSize;
    PIMAGE_IMPORT_BY_NAME imortByName;
    ULONG functionAddress;

    if (!GetImageDataDirectory(mNewImageBase, IMAGE_DIRECTORY_ENTRY_IMPORT, (PULONG) &imageImportDescriptor,
        &importSize))
    {
        KiDebug("GetModuleDataDirectory Failed");
        return FALSE;
    }

    while (imageImportDescriptor->OriginalFirstThunk && imageImportDescriptor->Name)
    {
        char moduleName[260];
        strcpy(moduleName, (char*) (mNewImageBase + imageImportDescriptor->Name));
        ULONG imageBase;
        if (!strcmp(_strlwr(mKernelName), _strlwr(moduleName)))
        {
            imageBase = mKernelImageBase;
            KiDebug("%s image base:0x%x", moduleName, imageBase);
        } else if (!GetImageBase(moduleName, imageBase))
        {
            KiDebug("FixKernelImportTable Failed");
            return FALSE;
        }

        if (imageBase == (ULONG) NULL )
        {
            KiDebug("FixModuleImportTable Alarm!");
            PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->FirstThunk);
            PIMAGE_THUNK_DATA OldFirstThunk = (PIMAGE_THUNK_DATA) ((ULONG) FirstThunk + mOldImageBase - mNewImageBase);
            while (OldFirstThunk->u1.Function)
            {
                FirstThunk->u1.Function = OldFirstThunk->u1.Function;
                OldFirstThunk++;
                FirstThunk++;
            }
            imageImportDescriptor++;
            continue;
        }
        PIMAGE_THUNK_DATA ImageThunkData = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->OriginalFirstThunk);
        PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA) (mNewImageBase + imageImportDescriptor->FirstThunk);
        KiDebug("FirstThunk address is:0x%x", FirstThunk);
        while (ImageThunkData->u1.Ordinal)
        {
            //序号导入
            if (IMAGE_SNAP_BY_ORDINAL32(ImageThunkData->u1.Ordinal))
            {
                if (!GetExportedRoutineByOrdinal(imageBase, ImageThunkData->u1.Ordinal & ~IMAGE_ORDINAL_FLAG32,
                    &functionAddress))
                {
                    KiDebug("GetExportedRoutineByOrdinal Failed");
                    return FALSE;
                }

                if (functionAddress == 0)
                {
                    KiDebug("can't find funcion Index %d \n",ImageThunkData->u1.Ordinal & ~IMAGE_ORDINAL_FLAG32);
                    break;
                }
                FirstThunk->u1.Function = functionAddress;
            }
            //函数名导入
            else
            {
                imortByName = (PIMAGE_IMPORT_BY_NAME) (mNewImageBase + ImageThunkData->u1.AddressOfData);
                if (!GetExportedRoutineByName(imageBase, (PCHAR)imortByName->Name, &functionAddress))
                {
                    KiDebug("GetExportedRoutineByName Failed");
                    return FALSE;
                }

                if (functionAddress == 0)
                {
                    KiDebug("can't Funcion Name:%s\n",imortByName->Name);
                    break;
                }
                FirstThunk->u1.Function = functionAddress;
            }
            FirstThunk++;
            ImageThunkData++;
        }
        imageImportDescriptor++;
    }
    return TRUE;
}

